import asyncio
import datetime
from typing import List

from llama_index.core.agent.workflow import  FunctionAgent
from llama_index.core.base.llms.types import ChatMessage
from llama_index.core.extractors import KeywordExtractor
from llama_index.core.storage.chat_store.sql import SQLAlchemyChatStore
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.schema import TextNode, NodeWithScore, QueryBundle
from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex, get_response_synthesizer
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.core.schema import Document
from pydantic import BaseModel

embed_model = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embed_model

from llama_index.llms.deepseek import DeepSeek

llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm
from llama_index.core.question_gen import LLMQuestionGenerator
from llama_index.question_gen.openai import OpenAIQuestionGenerator
from llama_index.llms.openai import OpenAI
from llama_index.core.question_gen.llm_generators import LLMQuestionGenerator
from llama_index.core.question_gen.output_parser import SubQuestionOutputParser


question_gen = LLMQuestionGenerator.from_defaults()
tools = [
    ToolMetadata(description="关于A相关的内容", name="source_1"),
    ToolMetadata(description="关于B相关的内容", name="source_2"),
]
query = QueryBundle(query_str="A和B的相同点有哪些?")
sub_questions = question_gen.generate(tools=tools, query=query)

print(sub_questions)